# Kestrel: An XMPP-based Many-Task Computing Scheduler
# Author: Lance Stout <lancestout@gmail.com>
#
# Copyright 2009-2010 Clemson University
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import threading

from task import Task


class Job(object):
    """
    Data structure representing a job.

    A job is an instance of a request to execute a command
    across various computing agents.
    """
    def __init__(self, data=None, state=[0]):
        """
        Instantiate a job object given a dictionary
        representation of a task that can include:

        source (string):
          The name of the user requesting tje job.
        job_id (int):
          The ID assigned to the job by the manager.
        command (list|string):
          If a string, the command is the path to an
          executable file. If it is a list, then each
          item in the list is an executable and will
          be executed in the order given.
        cleanup (string):
          An optional path to an executable responsible
          for cleaning up and resetting any resources used.
          If the agent processing the task goes offline
          while executing the cleanup command, the task
          is still considered completed.
        queue (int):
          The queue value is the number of tasks that
          are requested for execution (i.e., the job
          will be executed N times where N=queue).
        requires (list):
          A list of string tags which must be present
          in a worker's profile for a task from this job
          to be accepted.

        The state parameter is intended for internal
        use only. It contains the value for the next
        job id that will be issued. Since the initial
        list assigned to state is created only at compile
        time, this is essentially a static variable in
        the Java sense.
        """
        if data is None:
            data = {}

        # The source and job ID provide a globally
        # unique ID for the task.
        self.source = data.get('source', '')

        # If a job id was provided in the job data,
        # use it. If not, assign the next available
        # job id.
        if data.get('job_id', None) is not None:
            self.job_id = int(data['job_id'])
        else:
            self.job_id = state[0]
            state[0] += 1

        # Executing the main command is the
        # purpose for the job.
        self.command = data.get('command', '')

        # Once the main command has completed,
        # the cleanup command is executed. If
        # a task is terminated by an agent going
        # offline while in the cleanup stage, the
        # task will be considered completed and
        # will not be rescheduled.
        self.cleanup = data.get('cleanup', '')

        # The number of tasks that will be issued for
        # this job. For each task, the main command and
        # the cleanup command will be executed.
        self.queue = int(data.get('queue', 1))

        # Tags which a worker must match to run tasks
        # from this job.
        self.requires = set(data.get('requires', []))

        self.tasks = {'queued': [],
                      'pending': [],
                      'running': [],
                      'complete': []}
        self.tasks['queued'] = range(self.queue)
        self.lock = threading.Lock()

    def has_tasks(self):
        """
        Check that the job has tasks remaining.
        """
        with self.lock:
            return len(self.tasks['queued']) > 0

    def is_completed(self):
        """
        Check that all tasks for the job have been executed.
        """
        with self.lock:
            return len(self.tasks['complete']) == self.queue

    def next_task(self):
        """
        Return the next task from the job's queue.
        """
        with self.lock:
            if len(self.tasks['queued']) == 0:
                if len(self.tasks['pending']):
                    task_id = self.tasks['pending'].pop()
                    self.tasks['queued'].append(task_id)
                else:
                    return None

            task = self.tasks['queued'].pop()

            self.tasks['pending'].append(task)

            return Task({'job_id': self.job_id,
                         'task_id': task,
                         'command': self.command,
                         'cleanup': self.cleanup})

    def start_task(self, task_id):
        """Mark a task as having started."""
        with self.lock:
            if task_id in self.tasks['pending']:
                self.tasks['pending'].remove(task_id)
                self.tasks['running'].append(task_id)

    def complete_task(self, task_id):
        """Mark a task as finished."""
        with self.lock:
            self.tasks['running'].remove(task_id)
            self.tasks['complete'].append(task_id)

    def requeue_task(self, task_id):
        """Add a task back into the queue."""
        with self.lock:
            if task_id in self.tasks['pending']:
                self.tasks['pending'].remove(task_id)
                self.tasks['queued'].append(task_id)

            elif task_id in self.tasks['running']:
                self.tasks['running'].remove(task_id)
                self.tasks['queued'].append(task_id)

    def cancel(self):
        """
        Mark the job as canceled.

        Any queued tasks will be immediately completed,
        running jobs will be marked completed as the
        tasks exit and return.
        """
        with self.lock:
            tasks = self.tasks['queued']
            self.tasks['queued'] = []
            for task in tasks:
                self.tasks['complete'].append(task)

    def id(self):
        """
        Return the job's unique ID, which is composed of a
        manager name, and a job number.
        """
        return (self.source, self.job_id)

    def matches(self, worker_tags):
        """
        Determine if a worker can accept a task from this job.
        """
        return self.requires.issubset(worker_tags)

    def __repr__(self):
        return "(JobID: %d, Tasks: %d/%d/%d)" % (self.job_id,
                                                 len(self.tasks['running']),
                                                 len(self.tasks['complete']),
                                                 self.queue)
